#!/usr/bin/env ruby

# Copyright (c) 2010-2012 Microsoft Corp.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#	http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

require 'rubygems'
require 'ruby-aws'
require 'matrix'
require './include/array'

module MOS
class AssignmentSet
	attr_reader :stats, :assignment_list

	def initialize(assignment_list)
		@assignment_list = assignment_list
		rebuild_stats
	end

	def rebuild_stats
		init_stats

		build_stats(@stats_sentences, @algorithms, @sentences, @workers)
		build_stats(@stats_workers, @algorithms, @workers, @sentences)

		compute_algorithm_ci95
	end

	# 95th percentile from a t-distribution with n degrees of freedom
	# determined with MATLAB using:
	# >> n = (1:100)';
	# >> t = tinv(.5*(1 + .95), n);
	def t95(n)
		t = [12.7062, 4.3027, 3.1824, 2.7764, 2.5706, 2.4469, 2.3646, 2.3060, 2.2622, 2.2281,
		      2.2010, 2.1788, 2.1604, 2.1448, 2.1314, 2.1199, 2.1098, 2.1009, 2.0930, 2.0860,
		      2.0796, 2.0739, 2.0687, 2.0639, 2.0595, 2.0555, 2.0518, 2.0484, 2.0452, 2.0423,
		      2.0395, 2.0369, 2.0345, 2.0322, 2.0301, 2.0281, 2.0262, 2.0244, 2.0227, 2.0211,
		      2.0195, 2.0181, 2.0167, 2.0154, 2.0141, 2.0129, 2.0117, 2.0106, 2.0096, 2.0086,
		      2.0076, 2.0066, 2.0057, 2.0049, 2.0040, 2.0032, 2.0025, 2.0017, 2.0010, 2.0003,
		      1.9996, 1.9990, 1.9983, 1.9977, 1.9971, 1.9966, 1.9960, 1.9955, 1.9949, 1.9944,
		      1.9939, 1.9935, 1.9930, 1.9925, 1.9921, 1.9917, 1.9913, 1.9908, 1.9905, 1.9901,
		      1.9897, 1.9893, 1.9890, 1.9886, 1.9883, 1.9879, 1.9876, 1.9873, 1.9870, 1.9867,
		      1.9864, 1.9861, 1.9858, 1.9855, 1.9853, 1.9850, 1.9847, 1.9845, 1.9842, 1.9840];

		# if n is not on the table, use a Gaussian approximation (for n -> inf)
		return (n < t.length) ? t[n] : 1.96;
	end

	def stats_from_values(node)
		values = node[:values]
		stats  = (node[:stats] ||= {})

		stats[:count]    = values.length
		stats[:mean]     = values.mean
		stats[:min]      = values.min
		stats[:max]      = values.max
		stats[:var]      = values.var
		stats[:std_dev]  = values.std_dev
		stats[:kurtosis] = values.kurtosis

		if stats[:count] > 1
			t = t95(stats[:count] - 1)
			stats[:ci] = t * Math.sqrt(stats[:var] / stats[:count])
		end
	end

	def compute_algorithm_ci95
		@algorithms.each do |algorithm|
			next if !@stats_sentences[algorithm]

			t  = 0
			mi = []
			nj = []

			v_wu = []
			@sentences.each do |sentence|
				next if !@stats_sentences[algorithm][sentence]
				if @stats_sentences[algorithm][sentence][:stats][:count] >= 2
					v_wu << @stats_sentences[algorithm][sentence][:stats][:var]
				        mi   << @stats_sentences[algorithm][sentence][:stats][:count]
				end
			end
			v_wu = v_wu.mean

			v_su = []
			@workers.each do |worker|
				next if !@stats_workers[algorithm][worker]
				if @stats_workers[algorithm][worker][:stats][:count] >= 2
					v_su << @stats_workers[algorithm][worker][:stats][:var]
					nj   << @stats_workers[algorithm][worker][:stats][:count]
				end
			end
			v_su = v_su.mean

			v_swu = []
			@sentences.each do |sentence|
				next if !@stats_sentences[algorithm][sentence]
				@workers.each do |worker|
					next if !@stats_sentences[algorithm][sentence][worker]
					v_swu << @stats_sentences[algorithm][sentence][worker][:stats][:mean]
					t += 1
				end
			end
			v_swu = v_swu.var

			mi2 = mi.map { |v| v ** 2 }.sum
			nj2 = nj.map { |v| v ** 2 }.sum

			if v_su && v_wu
				m = Matrix[[ 1.0, 0.0, 1.0 ], [ 0.0, 1.0, 1.0 ], [ 1.0, 1.0, 1.0 ]]
				c = Matrix[[ v_su ], [ v_wu ], [ v_swu ]]
				v = m.inverse * c
				v_s = [ v[0,0], 0.0 ].max
				v_w = [ v[1,0], 0.0 ].max
				v_u = [ v[2,0], 0.0 ].max

				v_mu = v_s * mi2/(t ** 2) + v_w * nj2/(t ** 2) + v_u/t
			elsif !v_su && v_wu
				m = Matrix[[ 0.0, 1.0 ], [ 1.0, 1.0 ]]
				c = Matrix[[ v_wu ], [ v_swu ]]
				v = m.inverse * c
				v_s  = [ v[0,0], 0.0 ].max
				v_wu = [ v[1,0], 0.0 ].max

				v_mu = v_s * mi2/(t ** 2) + v_wu/t
			elsif v_su && !v_wu
				m = Matrix[[ 0.0, 1.0 ], [ 1.0, 1.0 ]]
				c = Matrix[[ v_su ], [ v_swu ]]
				v = m.inverse * c
				v_w  = [ v[0,0], 0.0 ].max
				v_su = [ v[1,0], 0.0 ].max

				v_mu = v_w * nj2/(t ** 2) + v_su/t
			else
				fail if v_su || v_wu
				v_mu = (v_swu == nil) ? 1e100 : v_swu/t
			end

			t = t95([ @sentences.length, @workers.length ].min - 1)
			@stats_sentences[algorithm][:stats][:ci] = t * Math.sqrt(v_mu)
			@stats_workers  [algorithm][:stats][:ci] = t * Math.sqrt(v_mu)
		end
	end

	def init_stats 
		@algorithms = {}
		@sentences  = {}
		@workers    = {}

		@stats_sentences = {}
		@stats_workers   = {}
		@assignment_list.each do |a|
			answer = a.answer
			worker = a.workerId.to_sym

			answer.each_key do |algorithm|
				answer[algorithm].each_key do |sentence|
					@algorithms[algorithm] = 1
					@sentences[sentence]   = 1
					@workers[worker]       = 1

					# initialize @stats_sentences
					@stats_sentences[:values] ||= []
					@stats_sentences[:stats ] ||= {}

					@stats_sentences[algorithm] ||= {}
					@stats_sentences[algorithm][:values] ||= []
					@stats_sentences[algorithm][:stats ] ||= {}

					@stats_sentences[algorithm][sentence] ||= {}
					@stats_sentences[algorithm][sentence][:values] ||= []
					@stats_sentences[algorithm][sentence][:stats ] ||= {}

					@stats_sentences[algorithm][sentence][worker] ||= {}
					@stats_sentences[algorithm][sentence][worker][:values] ||= []
					@stats_sentences[algorithm][sentence][worker][:stats ] ||= {}
					@stats_sentences[algorithm][sentence][worker][:values] << answer[algorithm][sentence]

					# initialize @stats_workers
					@stats_workers[:values] ||= []
					@stats_workers[:stats ] ||= {}

					@stats_workers[algorithm] ||= {}
					@stats_workers[algorithm][:values] ||= []
					@stats_workers[algorithm][:stats ] ||= {}

					@stats_workers[algorithm][worker] ||= {}
					@stats_workers[algorithm][worker][:values] ||= []
					@stats_workers[algorithm][worker][:stats ] ||= {}

					@stats_workers[algorithm][worker][sentence] ||= {}
					@stats_workers[algorithm][worker][sentence][:values] ||= []
					@stats_workers[algorithm][worker][sentence][:stats ] ||= {}
					@stats_workers[algorithm][worker][sentence][:values] << answer[algorithm][sentence]
				end
			end
		end

		@algorithms = @algorithms.keys.sort do |a,b|
			ref_a = a.to_s.match(/^Ref/i)
			ref_b = b.to_s.match(/^Ref/i)

			anc_a = a.to_s.match(/^Anc/i)
			anc_b = b.to_s.match(/^Anc/i)

			if ref_a && ref_b
				0
			elsif ref_a
				-1
			elsif ref_b
				1
			elsif anc_a && anc_b
				0
			elsif anc_a
				-1
			elsif anc_b
				1
			else
				a.to_s <=> b.to_s
			end
		end
		@sentences = @sentences.keys.sort { |a,b| a.to_s <=> b.to_s }
		@workers = @workers.keys.sort { |a,b| a.to_s <=> b.to_s }
	end

	def build_stats(stats, key1, key2, key3)
		key1.each do |k1|
			next if !stats[k1]

			key2.each do |k2|
				next if !stats[k1][k2]

				key3.each do |k3|
					next if !stats[k1][k2][k3]

					stats_from_values(stats[k1][k2][k3])
					stats[k1][k2][:values] << stats[k1][k2][k3][:stats][:mean]
					stats[k1][:values]     << stats[k1][k2][k3][:stats][:mean]
				end

				stats_from_values(stats[k1][k2])
#				stats[k1][:values] << stats[k1][k2][:stats][:mean]
			end

			stats_from_values(stats[k1])
		end
	
		return stats
	end

	def headphones(worker)
		assignment_list.each do |a|
			if a.workerId.to_sym == worker
				return a.headphones
			end
		end

		return ''
	end

	def compute_bonuses(min_assignments, min_working_time, bonus_quantity, bonus_quality_50pct, bonus_quality_10pct)
		bonus = {}
		del_list = {}

		@workers.each do |worker|
			bonus[worker] ||= {}
			bonus[worker][:amount] ||= 0.0
			bonus[worker][:reason] = ''
		end

		@workers.each do |worker|
			if assignment_count(worker) >= min_assignments
				bonus[worker][:amount] += bonus_quantity
				bonus[worker][:reason]  = "At least #{min_assignments} HITs completed"
			else
				del_list[worker] = 1
			end
		end

		@assignment_list.delete_if { |a| del_list[a.workerId.to_sym] }
		rebuild_stats

		correlation = []
		@workers.each do |worker|
			entry = {}
			entry[:worker] = worker
			entry[:correlation] = worker_correlation(worker)
			correlation << entry
		end
		correlation.sort! { |x,y| y[:correlation] <=> x[:correlation] }

		(0..((0.5 * correlation.length).round - 1)).each do |i|
			worker = correlation[i][:worker]
			bonus[worker][:amount] += bonus_quality_50pct
			bonus[worker][:reason]  = "At least #{min_assignments} HITs completed; set in the top 50%"
		end
		(0..((0.1 * correlation.length).round - 1)).each do |i|
			worker = correlation[i][:worker]
			bonus[worker][:amount] += bonus_quality_10pct
			bonus[worker][:reason]  = "At least #{min_assignments} HITs completed; set in the top 10%"
		end

		bonus
	end

	def fast_workers(min_working_time)
		worker_list = {}
		@assignment_list.each { |a| worker_list[a.workerId.to_sym] = 1 if a.workingTime < min_working_time }
		worker_list
	end

	def bad_workers(min_assignments)
		worker_list = {}
		@workers.each do |worker|
			if worker_correlation(worker) < 0.25 #&& assignment_count(worker) >= min_assignments
				worker_list[worker] = 1
			end
		end
		worker_list
	end

	def outliers(min_assignments)
		worker_list = {}
		@workers.each do |worker|
#			rw, rref = range(worker)
#			worker_list[worker] = 1 if (rref > 0.3) && (rw < 0.5 * rref)
			worker_list[worker] = 1 if worker_correlation(worker) < 0.7
			worker_list[worker] = 1 if assignment_count(worker) < min_assignments
		end
		worker_list
	end

	def print_worker_stats(workers = @workers)
		workers.each do |worker|
			printf("Stats for Worker %s\n", worker)
			printf("%-15s %-15s %-10s %-10s %-10s %-10s %-10s %-10s %-10s\n",
				'Sentence', 'Algorithm', 'Count', 'Min', 'Max', 'Mean', 'StdDev', 'Kurt', 'CI95')

			@algorithms.each do |algorithm|
				if !@stats_workers[algorithm] || !@stats_workers[algorithm][worker]
					next
				end

				s = @stats_workers[algorithm][worker][:stats]
				printf("%-15s %-15s %-10s %-10s %-10s %-10s %-10s %-10s %-10s\n",
					'--',
					algorithm,
					s[:count]    ? sprintf("%d",   s[:count]   ) : '---',
					s[:min]      ? sprintf("%.1f", s[:min]     ) : '---',
					s[:max]      ? sprintf("%.1f", s[:max]     ) : '---',
					s[:mean]     ? sprintf("%.2f", s[:mean]    ) : '----',
					s[:std_dev]  ? sprintf("%.2f", s[:std_dev] ) : '----',
					s[:kurtosis] ? sprintf("%.2f", s[:kurtosis]) : '----',
					s[:ci]       ? sprintf("%.2f", s[:ci]      ) : '----')
			end

			printf("\n")
			printf("%-15s %-15s %-10s %-10s\n",
				'Sentence', 'Algorithm', 'Count', 'Mean')

			@sentences.each do |sentence|
				@algorithms.each do |algorithm|
					if !@stats_workers[algorithm] ||
					   !@stats_workers[algorithm][worker] ||
					   !@stats_workers[algorithm][worker][sentence]
						next
					end

					s = @stats_workers[algorithm][worker][sentence][:stats]
					printf("%-15s %-15s %-10s %-10s\n",
						sentence,
						algorithm,
						s[:count] ? sprintf("%d",   s[:count]   ) : '---',
						s[:mean]  ? sprintf("%.2f", s[:mean]    ) : '----')
				end
			end

			printf("\n")
			printf("Audio setup: %s\n", headphones(worker))
			printf("Number of assignments: %d\n", assignment_count(worker))
			printf("Mean working time: %d\n", mean_working_time(worker))
			printf("Correlation: %4.2f\n", worker_correlation(worker))
			printf("\n")
		end
	end

	def print_sentence_stats
		printf("Global Stats\n")
		printf("%-15s %-15s %-10s %-10s %-10s %-10s %-10s %-10s %-10s\n",
			'Sentence', 'Algorithm', 'Count', 'Min', 'Max', 'Mean', 'StdDev', 'Kurt', 'CI95')

		@sentences.each do |sentence|
			@algorithms.each do |algorithm|
				if !@stats_sentences[algorithm] || !@stats_sentences[algorithm][sentence]
					next
				end

				s = @stats_sentences[algorithm][sentence][:stats]
				printf("%-15s %-15s %-10s %-10s %-10s %-10s %-10s %-10s %-10s\n",
					sentence,
					algorithm,
					s[:count]    ? sprintf("%d",   s[:count]   ) : '---',
					s[:min]      ? sprintf("%.1f", s[:min]     ) : '---',
					s[:max]      ? sprintf("%.1f", s[:max]     ) : '---',
					s[:mean]     ? sprintf("%.2f", s[:mean]    ) : '----',
					s[:std_dev]  ? sprintf("%.2f", s[:std_dev] ) : '----',
					s[:kurtosis] ? sprintf("%.2f", s[:kurtosis]) : '----',
					s[:ci]       ? sprintf("%.2f", s[:ci]      ) : '----')
			end
		end
		printf("\n")
	end

	def print_algorithm_stats
		printf("Global Stats, using %d assignments, %d subjects\n", assignment_list.length, @workers.length)

		printf("%-15s %-15s %-10s %-10s %-10s %-10s %-10s %-10s %-10s\n",
			'Sentence', 'Algorithm', 'Count', 'Min', 'Max', 'Mean', 'StdDev', 'Kurt', 'CI95')
		@algorithms.each do |algorithm|
			next if !@stats_sentences[algorithm]

			s = @stats_sentences[algorithm][:stats]
			printf("%-15s %-15s %-10s %-10s %-10s %-10s %-10s %-10s %-10s\n",
				'--',
				algorithm,
				s[:count]    ? sprintf("%d",   s[:count]   ) : '---',
				s[:min]      ? sprintf("%.1f", s[:min]     ) : '---',
				s[:max]      ? sprintf("%.1f", s[:max]     ) : '---',
				s[:mean]     ? sprintf("%.2f", s[:mean]    ) : '----',
				s[:std_dev]  ? sprintf("%.2f", s[:std_dev] ) : '----',
				s[:kurtosis] ? sprintf("%.2f", s[:kurtosis]) : '----',
				s[:ci]       ? sprintf("%.2f", s[:ci]      ) : '----')
		end
		printf("\n")
	end

	def assignment_count(worker)
		count = 0
		@assignment_list.each do |a|
			if a.workerId.to_sym == worker
				count += 1
			end
		end
		count
	end

	# consistency test according to ITU-R BT.500
	def consistent?(worker)
		p = 0
		q = 0
		n = 0
		@sentences.each do |sentence|
			@algorithms.each do |algorithm|
				if !@stats_sentences[algorithm] ||
				   !@stats_sentences[algorithm][sentence] ||
				   !@stats_sentences[algorithm][sentence][worker]
					next
				end

				s = @stats_sentences[algorithm][sentence][worker][:stats]
				r = @stats_sentences[algorithm][sentence][:stats]

				if (r[:kurtosis] - 3.0).abs <= 1.0
					if s[:mean] > r[:mean] + 2.0 * r[:std_dev]
						p += 1
					end
					if s[:mean] < r[:mean] - 2.0 * r[:std_dev]
						q += 1
					end
				else
					if s[:mean] > r[:mean] + Math.sqrt(20.0) * r[:std_dev]
						p += 1
					end
					if s[:mean] < r[:mean] - Math.sqrt(20.0) * r[:std_dev]
						q += 1
					end
				end

				n += 1
			end
		end

		if (p+q) > 0.0 && (p+q) / n > 0.05 && (p-q).abs / (p+q) < 0.3
			return false
		else
			return true
		end
	end

	def range(worker)
		v_mean = []
		v_work = []
		@algorithms.each do |algorithm|
			next if !@stats_workers[algorithm] || !@stats_workers[algorithm][worker]

			v_work << @stats_workers[algorithm][worker][:stats][:mean]
			v_mean << @stats_workers[algorithm][:stats][:mean]
		end

		work_range = v_work.max - v_work.min
		mean_range = v_mean.max - v_mean.min

		return work_range, mean_range
	end

	def worker_correlation(worker)
		a_correlation = algorithm_mos_correlation(worker)
		s_correlation = sentence_mos_correlation(worker)

		fail unless a_correlation || s_correlation
		return a_correlation if a_correlation
		return s_correlation if s_correlation
	end

	def sentence_mos_correlation(worker)
		v1 = []
		v2 = []

		@sentences.each do |sentence|
			@algorithms.each do |algorithm|
				if !@stats_sentences[algorithm] ||
				   !@stats_sentences[algorithm][sentence] ||
				   !@stats_sentences[algorithm][sentence][worker]
					next
				end

				v1 << @stats_sentences[algorithm][sentence][:stats][:mean]
				v2 << @stats_sentences[algorithm][sentence][worker][:stats][:mean]
			end
		end

		Array::correlation_coefficient(v1, v2)
	end

	# used for finding outliers and computing bonuses
	def algorithm_mos_correlation(worker)
		v1 = []
		v2 = []

		@algorithms.each do |algorithm|
			if !@stats_workers[algorithm] ||
			   !@stats_workers[algorithm][worker]
				next
			end

			v1 << @stats_workers[algorithm][:stats][:mean]
			v2 << @stats_workers[algorithm][worker][:stats][:mean]
		end

		Array::correlation_coefficient(v1, v2)
	end

	def mean_working_time(worker = nil)
		if worker
			@assignment_list.find_all {|a| a.workerId.to_sym == worker }.map { |a| a.workingTime }.mean
		else
			@assignment_list.map { |a| a.workingTime }.mean
		end
	end

	def print_raw_scores
		datafile = File.new("MATLAB/algorithms.txt", "w")
		datafile.puts @algorithms.join("\t")
		datafile.close

		datafile = File.new("MATLAB/sentences.txt", "w")
		datafile.puts @sentences.join("\t")
		datafile.close

		@algorithms.each do |algorithm|
			filename = sprintf("MATLAB/scores_%s.txt", algorithm)
			datafile = File.new(filename, "w")

			@workers.each do |worker|
				v = []
				@sentences.each do |sentence|
					if @stats_sentences[algorithm][sentence] &&
					   @stats_sentences[algorithm][sentence][worker] &&
					   @stats_sentences[algorithm][sentence][worker][:stats][:mean]
						v << @stats_sentences[algorithm][sentence][worker][:stats][:mean]
					else
						v << -1.0
					end
				end
				datafile.puts v.map { |score| sprintf('%.3f', score) }.join("\t")
			end

			datafile.close
		end
	end

end	# AssignmentSet
end	# MOS
